function [bpfa,VBparam] = BPFA(X,K,maxite)
% BPFA factorizes the matrix X with a linear Gaussian factorization. 
% It uses a finite approximation of the beta process as a prior.
%
% X: D x N matrix
% K: Number of atoms (truncation)
% maxite: Maximum number of iterations
%

    [D,N] = size(X);
    gamma = 1;
    alpha = 1;
    sigw2 = 1;
    sign2 = 1*var(X(:))/5;
    sigphi2 = 10^6;

    sigPhi = zeros(1,K);
    sigW = zeros(K,N);

    Phi = KmeansBPFA(X,K,30);  % DxK
    % [U,S,V] = svds(X,K);
    % Phi = U*S/sqrt(N); clear U S V;

    tic;
    fprintf('OMP begin.....');
    W = zeros(K,N);
    Z = zeros(K,N);
    for n = 1:N
        [W(:,n),Z(:,n)] = OMP(X(:,n), Phi, ceil(.25*D), (100/16)*10^-2);
    end
    endT1 = toc();
    fprintf('finished. Time cost %f seconds\n', endT1);

    a = ones(K,1);
    b = ones(K,1);

    bool = 1; 
    ite = 0; 
    Lbound_const = -.5*N*D*log(2*pi*sign2) + K*gammaln(alpha) - K*gammaln(alpha*gamma/K) - K*gammaln(alpha*(1-gamma/K));
    Lbound = [];

    fprintf('BPFA Trainning beging...\n');
    while bool
        tic;
        ite = ite + 1;
        Lbound(ite) = Lbound_const;

      % Update Phi
        for num = 1:3
            for k = 1:K
                idx = 1:K;
                idx(k) = [];
                RES = X - Phi(:,idx)*(W(idx,:).*Z(idx,:));
                vec = RES*(W(k,:).*Z(k,:))'/sign2;
                sigPhi(k) = 1/(1/sigphi2 + (W(k,:).^2+sigW(k,:))*Z(k,:)'/sign2);
                Phi(:,k) = vec*sigPhi(k);
            end
        end

      % Update W
        for n = 1:N
            PhiZ = Phi.*repmat(Z(:,n)',D,1);
            % ... change....
            invM =  sigw2*eye(K) - sigw2^2*PhiZ'*( ( sigw2*eye(D)+sign2*(PhiZ*PhiZ') )\PhiZ);
            %invM =  sigw2*eye(K) - sigw2^2*PhiZ'*inv( sign2*eye(D)+sigw2*(PhiZ*PhiZ') )*PhiZ;
            W(:,n) = invM*(PhiZ'*X(:,n))/sign2;
            sigW(:,n) = diag(invM);
        end

      % Update Z
        for k = 1:K
            idx = 1:K;
            idx(k) = [];
            RES = X - Phi(:,idx)*(W(idx,:).*Z(idx,:));
            BasisAdd = Phi(:,k)*W(k,:);
            lnphi_nk0 = zeros(1,N);
            lnphi_nk1 = zeros(1,N);
            if ite > 1
                lnphi_nk0 = lnphi_nk0 + psi(b(k)) - psi(a(k)+b(k));
                lnphi_nk1 = lnphi_nk1 + psi(a(k)) - psi(a(k)+b(k));
            end
            lnphi_nk1 = lnphi_nk1 - .5*(sum(BasisAdd.^2,1) - 2*sum(BasisAdd.*RES,1))/sign2;
            lnphi_nk1 = lnphi_nk1 - .5*(sigW(k,:)*(Phi(:,k)'*Phi(:,k) + D*sigPhi(k)) + D*sigPhi(k)*W(k,:).^2)/sign2;
            lnphi_nk = [lnphi_nk1 ; lnphi_nk0];
            lnphi_nk = lnphi_nk - repmat(max(lnphi_nk,[],1),2,1);
            phi_nk = exp(lnphi_nk);
            Z(k,:) = round(phi_nk(1,:)./sum(phi_nk,1)); % Indicator function q distribution on Z
    %         Z(k,:) = phi_nk(1,:)./sum(phi_nk,1);
        end
        n_k = sum(Z,2);
        a = n_k + alpha*gamma/K;
        b = N - n_k + alpha*(1-gamma/K);

      % Update Lower Bound
        vec = sum(Phi.^2,1) + D*sigPhi;
        for n = 1:N
            Lbound(ite) = Lbound(ite) - .5*(sum((X(:,n)-Phi*(W(:,n).*Z(:,n))).^2) + vec*(Z(:,n).*sigW(:,n)) + D*sigPhi*(Z(:,n).*(W(:,n).^2+sigW(:,n))))/sign2;
        end
        Lbound(ite) = Lbound(ite) - .5*sum(W(:).^2 + sigW(:))/sigw2 + .5*sum(log(sigW(:)/sigw2)) - .5*sum(vec)/sigphi2 + D*.5*sum(log(sigPhi)/sigphi2);
        Lbound(ite) = Lbound(ite) - sum(gammaln(a+b)) + sum(gammaln(a)) + sum(gammaln(b)); % - (a-1)'*(psi(a)-psi(a+b)) - (b-1)'*(psi(b)-psi(a+b));

        if ite > 1
            if abs(Lbound(ite)-Lbound(ite-1))/abs(Lbound(ite-1)) < 10^-3
                bool = 0;
            elseif ite == maxite
                bool = 0;
            end
        end

        endT2 = toc();
        fprintf('BPFA:Iteration:%d, Time cost: %f, lowbound:%f \n ', ite, endT2, Lbound(ite));
    end

    %imagesc(Z); title(num2str(ite)); pause(.1);
    %VBpoint = {Lbound,time1,time2};
    VBparam = {Phi,W,Z,a,b};
    
    bpfa.K = K;
    bpfa.Phi = Phi;

end